from ultralytics import YOLO
from PIL import Image
import torch
import cv2

class All_Classify():
    def __init__(self):
        self.C_model = YOLO(r'./AI_weights\best_compoment_cls_0119.pt')

    def start_classify(self,datadict):
        res_dict = {}
        for key,image in datadict.items():
            # cv2.imshow('a', image)
            # cv2.waitKey(0)
            image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))

            res = self.C_model(image)
            classes = res[0].names
            class_name = classes[res[0].probs.top1]
            conf = res[0].probs.top1conf
            if conf >=0:

                res_dict[key] = class_name
            else:
                res_dict[key] = 'Other'

        return res_dict
    def run_img(self,image):
        image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        res = self.C_model(image)
        classes = res[0].names
        class_name = classes[res[0].probs.top1]
        conf = res[0].probs.top1conf
        return class_name,conf

class C_Classify():
    def __init__(self):
        self.C_model = YOLO(r'./AI_weights/C_best_0118.pt')

    def start_classify(self,datalist,ng_dict):
        res_dict = {}
        for data in datalist:
            key,image= data[0],data[2]
            # cv2.imshow('a',image)
            # cv2.waitKey(0)
            image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
            res = self.C_model(image)
            classes = res[0].names
            class_name = classes[res[0].probs.top1]
            if class_name in ['1','2','3','4','5','6','7','8','9','10','11','12']:
                continue
            else:
                ng_dict[key] = class_name

        return ng_dict

    def run_img(self,image):
        image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        res = self.C_model(image)
        classes = res[0].names
        class_name = classes[res[0].probs.top1]
        conf = res[0].probs.top1conf
        return class_name,conf

class R_Classify():
    def __init__(self):
        self.R_model = YOLO(r'./AI_weights\R_best_0118.pt')

    def start_classify(self,datalist,ng_dict):
        res_dict = {}
        for data in datalist:
            key,image= data[0],data[2]
            # cv2.imshow('a', image)
            # cv2.waitKey(0)
            image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
            res = self.R_model(image)
            classes = res[0].names
            class_name = classes[res[0].probs.top1]
            if class_name in ['R1','R2','R3','R4','R5','R6']:
                continue
            else:
                ng_dict[key] = class_name

        return ng_dict

    def run_img(self,image):
        image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        res = self.R_model(image)
        classes = res[0].names
        class_name = classes[res[0].probs.top1]
        conf = res[0].probs.top1conf
        return class_name,conf


class D_Classify():
    def __init__(self):
        self.D_model = YOLO(r'./AI_weights/D_best_1127.pt')

    def start_classify(self,datalist,ng_dict):
        res_dict = {}
        for data in datalist:
            key,image= data[0],data[2]
            # cv2.imshow('a',image)
            # cv2.waitKey(0)
            image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
            res = self.D_model(image)
            classes = res[0].names
            class_name = classes[res[0].probs.top1]
            if class_name in ['D1','D2','D3','D4','D5']:
                continue
            else:
                ng_dict[key] = class_name

        return ng_dict

    def run_img(self,image):
        image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        res = self.D_model(image)
        classes = res[0].names
        class_name = classes[res[0].probs.top1]
        conf = res[0].probs.top1conf
        return class_name,conf

class Q_Classify():
    def __init__(self):
        self.Q_model = YOLO(r'./AI_weights/Q_best_1127.pt')

    def start_classify(self,datalist,ng_dict):
        res_dict = {}
        for data in datalist:
            key,image= data[0],data[2]
            # cv2.imshow('a',image)
            # cv2.waitKey(0)
            image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
            res = self.Q_model(image)
            classes = res[0].names
            class_name = classes[res[0].probs.top1]
            if class_name in ['Q1','Q2','Q3','Q4']:
                continue
            else:
                ng_dict[key] = class_name

        return ng_dict

    def run_img(self,image):
        image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        res = self.Q_model(image)
        classes = res[0].names
        class_name = classes[res[0].probs.top1]
        conf = res[0].probs.top1conf
        return class_name,conf
